using System;
using System.Collections.Generic;
using System.Runtime.CompilerServices;
using UnityEngine;

namespace HTCollections.Custom
{
    /// <summary>
    /// 二维KD树
    /// </summary>
    public sealed class KdTree2D
    {
        /// <summary>
        /// 最大叶子节点大小
        /// </summary>
        public int MaxLeafSize;

        public Node[] Nodes;
        Vector3[] mPoints;

        public KdTree2D(int maxLeafSize = 5)
        {
            MaxLeafSize = maxLeafSize;
        }

        /// <summary>
        /// 构建KD树
        /// </summary>
        /// <param name="points">坐标点集合</param>
        public void Build(Vector3[] points)
        {
            if (points == null || points.Length == 0) return;
            mPoints = points; //引用传递

            if (Nodes == null || Nodes.Length != points.Length * 2)
                Nodes = new Node[points.Length * 2];

            BuildRecursive(0, points.Length - 1, 0);
        }

        void BuildRecursive(int begin, int end, int nodeIndex)
        {
            Nodes[nodeIndex] = new Node(nodeIndex, begin, end, float.MaxValue, float.MinValue, float.MaxValue,
                float.MinValue, -1, -1);

            //计算当前节点的最小和最大坐标
            for (int i = begin; i <= end; i++)
            {
                Nodes[nodeIndex].MinX = Min(Nodes[nodeIndex].MinX, mPoints[i].x);
                Nodes[nodeIndex].MaxX = Max(Nodes[nodeIndex].MaxX, mPoints[i].x);
                Nodes[nodeIndex].MinY = Min(Nodes[nodeIndex].MinY, mPoints[i].y);
                Nodes[nodeIndex].MaxY = Max(Nodes[nodeIndex].MaxY, mPoints[i].y);
            }

            //如果当前节点的代理数量小于等于最大叶子节点大小，则不需要继续递归构建子树
            if (end - begin + 1 <= MaxLeafSize) return;

            int leftEnd = begin - 1;
            //计算当前节点的分割方向和分割值
            if (Nodes[nodeIndex].MaxX - Nodes[nodeIndex].MinX >
                Nodes[nodeIndex].MaxY - Nodes[nodeIndex].MinY) //沿竖直方向进行分割
            {
                float splitValue = (Nodes[nodeIndex].MaxX + Nodes[nodeIndex].MinX) / 2;
                //划分代理到左右子树
                for (int i = begin; i <= end; i++)
                {
                    if (mPoints[i].x <= splitValue) //左子树
                    {
                        leftEnd++;
                        Swap(i, leftEnd);
                    }
                }
            }
            else //沿水平方向进行分割
            {
                float splitValue = (Nodes[nodeIndex].MaxY + Nodes[nodeIndex].MinY) / 2;
                //划分代理到左右子树
                for (int i = begin; i <= end; i++)
                {
                    if (mPoints[i].y <= splitValue) //左子树
                    {
                        leftEnd++;
                        Swap(i, leftEnd);
                    }
                }
            }

            //如果左子树或右子树为空，则不需要继续递归构建子树
            if (leftEnd < begin || leftEnd >= end) return;

            //递归构建左子树和右子树
            Nodes[nodeIndex].LeftIndex = nodeIndex + 1;
            Nodes[nodeIndex].RightIndex = nodeIndex + 2 * (leftEnd - begin + 1);
            BuildRecursive(begin, leftEnd, Nodes[nodeIndex].LeftIndex);
            BuildRecursive(leftEnd + 1, end, Nodes[nodeIndex].RightIndex);
        }

        /// <summary>
        /// 搜索邻近点
        /// </summary>
        /// <param name="point">搜索中心位置坐标</param>
        /// <param name="range">搜索半径</param>
        /// <param name="results">搜索结果</param>
        public void QueryNeighborList(Vector3 point, float range, ref List<Vector3> results)
        {
            QueryNeighborList(point, range * range, 0, ref results);
        }

        void QueryNeighborList(Vector3 point, float rangeSq, int nodeIndex, ref List<Vector3> results)
        {
            //如果当前节点是叶子节点
            if (Nodes[nodeIndex].End - Nodes[nodeIndex].Begin + 1 <= MaxLeafSize)
            {
                for (int i = Nodes[nodeIndex].Begin; i <= Nodes[nodeIndex].End; i++)
                {
                    float disSq = DistanceSqrt(point, mPoints[i]);
                    if (disSq < rangeSq)
                        results.Add(mPoints[i]);
                }
            }
            //如果当前节点不是叶子节点
            else
            {
                //计算指定代理到当前节点左子树和右子树的最短距离的平方
                //代理在左子树、右子树内部时最短距离视为0
                float disSq_left =
                    Sqrt(Max(0f, Nodes[Nodes[nodeIndex].LeftIndex].MinX - point.x)) +
                    Sqrt(Max(0f, point.x - Nodes[Nodes[nodeIndex].LeftIndex].MaxX)) +
                    Sqrt(Max(0f, Nodes[Nodes[nodeIndex].LeftIndex].MinY - point.y)) +
                    Sqrt(Max(0f, point.y - Nodes[Nodes[nodeIndex].LeftIndex].MaxY));
                float disSq_right =
                    Sqrt(Max(0f, Nodes[Nodes[nodeIndex].RightIndex].MinX - point.x)) +
                    Sqrt(Max(0f, point.x - Nodes[Nodes[nodeIndex].RightIndex].MaxX)) +
                    Sqrt(Max(0f, Nodes[Nodes[nodeIndex].RightIndex].MinY - point.y)) +
                    Sqrt(Max(0f, point.y - Nodes[Nodes[nodeIndex].RightIndex].MaxY));

                if (disSq_left < disSq_right) //如果左子树更近
                {
                    if (disSq_left < rangeSq) //如果左子树可能存在距离指定代理在一定范围内的邻居代理
                    {
                        QueryNeighborList(point, rangeSq, Nodes[nodeIndex].LeftIndex, ref results); //递归查询左子树
                        if (disSq_right < rangeSq) //如果右子树也可能存在距离指定代理在一定范围内的邻居代理
                            QueryNeighborList(point, rangeSq, Nodes[nodeIndex].RightIndex, ref results); //递归查询右子树
                    }
                }
                else //如果右子树更近
                {
                    if (disSq_right < rangeSq) //如果右子树可能存在距离指定代理在一定范围内的邻居代理
                    {
                        QueryNeighborList(point, rangeSq, Nodes[nodeIndex].RightIndex, ref results); //递归查询右子树
                        if (disSq_left < rangeSq) //如果左子树可能存在距离指定代理在一定范围内的邻居代理
                            QueryNeighborList(point, rangeSq, Nodes[nodeIndex].LeftIndex, ref results); //递归查询左子树
                    }
                }
            }
        }

        /// <summary>
        /// 搜索指定矩形范围内的点
        /// </summary>
        /// <param name="rect">指定矩形范围（minX, minY, maxX, maxY）</param>
        /// <param name="results">结果集合</param>
        public void QueryNeighborList(Vector4 rect, ref List<Vector3> results)
        {
            QueryNeighborList(rect, 0, ref results);
        }

        void QueryNeighborList(Vector4 rect, int nodeIndex, ref List<Vector3> results)
        {
            //如果当前节点是叶子节点
            if (Nodes[nodeIndex].End - Nodes[nodeIndex].Begin + 1 <= MaxLeafSize)
            {
                for (int i = Nodes[nodeIndex].Begin; i <= Nodes[nodeIndex].End; i++)
                {
                    if (Contain(mPoints[i], rect))
                        results.Add(mPoints[i]);
                }
            }
            //如果当前节点不是叶子节点
            else
            {
                //如果左子树可能存在矩形区域内的代理
                if (Contain(rect, Nodes[Nodes[nodeIndex].LeftIndex]))
                {
                    QueryNeighborList(rect, Nodes[nodeIndex].LeftIndex, ref results); //递归查询左子树
                }

                //如果右子树可能存在矩形区域内的代理
                if (Contain(rect, Nodes[Nodes[nodeIndex].RightIndex]))
                {
                    QueryNeighborList(rect, Nodes[nodeIndex].RightIndex, ref results); //递归查询右子树
                }
            }
        }

        /// <summary>
        /// 搜索指定范围内的最邻近点
        /// </summary>
        /// <param name="point">搜索中心位置坐标</param>
        /// <param name="range">搜索半径</param>
        /// <returns>坐标索引</returns>
        public int QueryNearest(Vector3 point, float range)
        {
            float rangeSq = float.MaxValue;
            int index = -1;
            QueryNearest(point, ref rangeSq, 0, ref index);
            if (rangeSq >= range * range) return -1;
            return index;
        }

        void QueryNearest(Vector3 point, ref float rangeSq, int nodeIndex, ref int index)
        {
            //如果当前节点是叶子节点
            if (Nodes[nodeIndex].End - Nodes[nodeIndex].Begin + 1 <= MaxLeafSize)
            {
                for (int i = Nodes[nodeIndex].Begin; i <= Nodes[nodeIndex].End; i++)
                {
                    float disSq = DistanceSqrt(point, mPoints[i]);
                    if (disSq < rangeSq)
                    {
                        rangeSq = disSq;
                        index = i;
                    }
                }
            }
            //如果当前节点不是叶子节点
            else
            {
                //计算指定代理到当前节点左子树和右子树的最短距离的平方
                //代理在左子树、右子树内部时最短距离视为0
                float disSq_left =
                    Sqrt(Max(0f, Nodes[Nodes[nodeIndex].LeftIndex].MinX - point.x)) +
                    Sqrt(Max(0f, point.x - Nodes[Nodes[nodeIndex].LeftIndex].MaxX)) +
                    Sqrt(Max(0f, Nodes[Nodes[nodeIndex].LeftIndex].MinY - point.y)) +
                    Sqrt(Max(0f, point.y - Nodes[Nodes[nodeIndex].LeftIndex].MaxY));
                float disSq_right =
                    Sqrt(Max(0f, Nodes[Nodes[nodeIndex].RightIndex].MinX - point.x)) +
                    Sqrt(Max(0f, point.x - Nodes[Nodes[nodeIndex].RightIndex].MaxX)) +
                    Sqrt(Max(0f, Nodes[Nodes[nodeIndex].RightIndex].MinY - point.y)) +
                    Sqrt(Max(0f, point.y - Nodes[Nodes[nodeIndex].RightIndex].MaxY));

                if (disSq_left < disSq_right) //如果左子树更近
                {
                    if (disSq_left < rangeSq) //如果左子树可能存在距离指定代理在一定范围内的邻居代理
                    {
                        QueryNearest(point, ref rangeSq, Nodes[nodeIndex].LeftIndex, ref index); //递归查询左子树
                        if (disSq_right < rangeSq) //如果右子树也可能存在距离指定代理在一定范围内的邻居代理
                            QueryNearest(point, ref rangeSq, Nodes[nodeIndex].RightIndex, ref index); //递归查询右子树
                    }
                }
                else //如果右子树更近
                {
                    if (disSq_right < rangeSq) //如果右子树可能存在距离指定代理在一定范围内的邻居代理
                    {
                        QueryNearest(point, ref rangeSq, Nodes[nodeIndex].RightIndex, ref index); //递归查询右子树
                        if (disSq_left < rangeSq) //如果左子树可能存在距离指定代理在一定范围内的邻居代理
                            QueryNearest(point, ref rangeSq, Nodes[nodeIndex].LeftIndex, ref index); //递归查询左子树
                    }
                }
            }
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        void Swap(int a, int b)
        {
            var temp = mPoints[a];
            mPoints[a] = mPoints[b];
            mPoints[b] = temp;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        float Sqrt(float a)
        {
            return a * a;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        float DistanceSqrt(Vector3 a, Vector3 b)
        {
            float x = a.x - b.x;
            float y = a.y - b.y;
            return x * x + y * y;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        float Min(float a, float b)
        {
            if (a > b) return b;
            return a;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        float Max(float a, float b)
        {
            if (a > b) return a;
            return b;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        bool Contain(Vector3 point, Vector4 rect)
        {
            return point.x >= rect.x && point.x <= rect.z && point.y >= rect.y && point.y <= rect.w;
        }

        [MethodImpl(MethodImplOptions.AggressiveInlining)]
        bool Contain(Vector4 rect, Node node)
        {
            if (node.MaxX < rect.x || node.MinX > rect.z) return false;
            if (node.MaxY < rect.y || node.MinY > rect.w) return false;
            return true;
        }

        public struct Node : IEquatable<Node>
        {
            public readonly int Index; //节点Id
            public int Begin; //起始索引
            public int End; //终止索引
            public float MinX; //x轴上的最小值
            public float MaxX; //x轴上的最大值
            public float MinY; //y轴上的最小值
            public float MaxY; //y轴上的最大值
            public int LeftIndex; //左子树节点索引
            public int RightIndex; //右子树节点索引

            public Node(int index, int begin, int end, float minX, float maxX, float minY, float maxY, int leftIndex,
                int rightIndex)
            {
                Index = index;
                Begin = begin;
                End = end;
                MinX = minX;
                MaxX = maxX;
                MinY = minY;
                MaxY = maxY;
                LeftIndex = leftIndex;
                RightIndex = rightIndex;
            }

            public bool Equals(Node other)
            {
                return Index == other.Index && Begin == other.Begin && End == other.End && MinX.Equals(other.MinX) &&
                       MaxX.Equals(other.MaxX) && MinY.Equals(other.MinY) && MaxY.Equals(other.MaxY) &&
                       LeftIndex == other.LeftIndex && RightIndex == other.RightIndex;
            }

            public override bool Equals(object obj)
            {
                return obj is Node other && Equals(other);
            }

            public override int GetHashCode()
            {
                return Index.GetHashCode();
            }
        }
    }
}